#pragma once
#include <cstdint>
#include <algorithm>
#include <iostream>
#include "./utils/random_mapping.hpp"


namespace clay {
    template<std::uint16_t charset_size>
    class Rotor final {
    private:
        std::uint8_t current;
        std::array<std::uint8_t, charset_size> mapping;
        std::array<std::uint8_t, charset_size> reverse_mapping;

        void buildReverseMapping(const std::array<std::uint8_t, charset_size>& mapping)
        {
            uint16_t i, map_to;
            for (i = 0; i < charset_size; i++) {
                map_to = mapping[i];
                reverse_mapping[map_to] = i;
            }
        }


        std::uint16_t roll(const std::uint16_t times)
        {
            std::uint16_t real_roll_times = times % charset_size;

            std::rotate(mapping.rbegin(), mapping.rbegin() + real_roll_times, mapping.rend());
            std::for_each(mapping.rbegin(), mapping.rend(),
                [&real_roll_times](std::uint8_t& value) {
                    std::uint16_t current_value = value;
                    value = (current_value + real_roll_times) % charset_size;
                }
            );

            std::rotate(reverse_mapping.rbegin(), reverse_mapping.rbegin() + real_roll_times, reverse_mapping.rend());
            std::for_each(reverse_mapping.rbegin(), reverse_mapping.rend(),
                [&real_roll_times](std::uint8_t& value) {
                    std::uint16_t current_value = value;
                    value = (current_value + real_roll_times) % charset_size;
                }
            );

            return real_roll_times;
        }
    public:
        Rotor() = default;
        Rotor(const std::uint8_t shift, const std::uint64_t seed)
            : current(0),
              mapping(generateRandomMapping<charset_size>(seed))
        {
            buildReverseMapping(this->mapping);
            roll(shift % charset_size);
        }


        std::uint8_t Roll(const std::uint16_t times = 1)
        {
            std::uint8_t carry = 0;
            std::uint16_t next_current = current + roll(times);

            current = next_current % charset_size;
            carry = next_current / charset_size;

            return carry;
        }


        std::uint8_t Map(const std::uint8_t value) const
        {
            return mapping.at(value);
        }


        std::uint8_t ReverseMap(const std::uint8_t value) const
        {
            return reverse_mapping.at(value);
        }


        void PrintMapping() const
        {
            std::cout << "mapping: ";
            for (int i = 0; i < charset_size; i++) {
                std::cout << i << "->" << static_cast<int>(mapping[i]) << " ";
            }
            std::cout << std::endl;

            std::cout << "rev_map: ";
            for (int i = 0; i < charset_size; i++) {
                std::cout << i << "->" << static_cast<int>(reverse_mapping[i]) << " ";
            }
            std::cout << std::endl;
        }
    };
}